//! Go2 Base Node
//! 
//! This node combines functionality from cmd_vel_demo and odom_demo to provide
//! a base controller for the Unitree Go2 robot that can receive velocity commands
//! and publish odometry information.

use r2r::{
    self, QosProfile,
    geometry_msgs::msg::{Twist, TransformStamped, PoseStamped},
    nav_msgs::msg::Odometry,
    tf2_msgs::msg::TFMessage,
    unitree_api::msg::Request,
};
use serde_json::json;
use std::sync::atomic::{AtomicBool, Ordering};
use futures::{executor::LocalPool, stream::StreamExt, task::LocalSpawnExt};
use std::time::{SystemTime, UNIX_EPOCH};

// Constants
static RUNNING: AtomicBool = AtomicBool::new(true);

/// Frame IDs used in the odometry system
const ODOM_FRAME: &str = "odom";
const BASE_LINK_FRAME: &str = "base_link";

/// Topic names
const CMD_VEL_TOPIC: &str = "/cmd_vel";
const POSE_TOPIC: &str = "/utlidar/robot_pose";
const ODOM_TOPIC: &str = "/odom";
const TF_TOPIC: &str = "/tf";
const SPORT_REQUEST_TOPIC: &str = "/api/sport/request";

/// Node name
const NODE_NAME: &str = "go2_base";

/// Safety limits
const MAX_LINEAR_X: f32 = 0.8;    // m/s
const MAX_LINEAR_Y: f32 = 0.4;    // m/s
const MAX_ANGULAR_VEL: f32 = 1.0; // rad/s

/// Get current system time as seconds since Unix epoch
/// 
/// # Returns
/// Current time as f64 seconds, or 0.0 if system time is invalid
fn get_current_time() -> f64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_secs_f64())
        .unwrap_or_else(|_| {
            eprintln!("Warning: System time is before Unix epoch, using 0.0");
            0.0
        })
}

/// Convert f64 timestamp to ROS2 Time message
/// 
/// # Arguments
/// * `timestamp` - Time in seconds since Unix epoch
/// 
/// # Returns
/// Tuple of (seconds, nanoseconds) for ROS2 Time
fn timestamp_to_ros_time(timestamp: f64) -> (i32, u32) {
    let secs = timestamp.floor() as i32;
    let nsecs = ((timestamp - secs as f64) * 1e9).round() as u32;
    (secs, nsecs)
}

/// Handle incoming cmd_vel messages and send move commands to robot
fn handle_cmd_vel(twist_msg: &Twist, req_puber: &r2r::Publisher<Request>) {
    // Extract and limit velocities
    let vx = (twist_msg.linear.x as f32).clamp(-MAX_LINEAR_X, MAX_LINEAR_X);
    let vy = (twist_msg.linear.y as f32).clamp(-MAX_LINEAR_Y, MAX_LINEAR_Y);
    let vyaw = (twist_msg.angular.z as f32).clamp(-MAX_ANGULAR_VEL, MAX_ANGULAR_VEL);
    
    // Log command
    if vx != 0.0 || vy != 0.0 || vyaw != 0.0 {
        println!("Move: vx={:.2}, vy={:.2}, vyaw={:.2}", vx, vy, vyaw);
    }
    
    // Send move command
    send_move_command(vx, vy, vyaw, req_puber);
}

/// Send move command to robot
fn send_move_command(vx: f32, vy: f32, vyaw: f32, req_puber: &r2r::Publisher<Request>) {
    let mut req = Request::default();
    
    // Create move command
    let js = json!({"x": vx, "y": vy, "z": vyaw});
    req.parameter = js.to_string();
    req.header.identity.api_id = 1008; // MOVE command
    
    // Publish
    if let Err(e) = req_puber.publish(&req) {
        eprintln!("Failed to send move command: {:?}", e);
    }
}

/// Send stop command to robot
fn send_stop_command(req_puber: &r2r::Publisher<Request>) {
    println!("Sending stop command...");
    send_move_command(0.0, 0.0, 0.0, req_puber);
}

/// Handle incoming pose messages and publish odometry and TF
fn handle_pose(
    msg: PoseStamped, 
    odom_puber: &r2r::Publisher<Odometry>, 
    tf_puber: &r2r::Publisher<TFMessage>
) -> Result<(), Box<dyn std::error::Error>> {
    let current_time = get_current_time();
    let (secs, nsecs) = timestamp_to_ros_time(current_time);
    
    // Create and publish messages directly to avoid cloning
    publish_tf_and_odometry(tf_puber, odom_puber, msg, secs, nsecs)
}

/// Publish both TF and odometry messages efficiently without unnecessary cloning
fn publish_tf_and_odometry(
    tf_puber: &r2r::Publisher<TFMessage>,
    odom_puber: &r2r::Publisher<Odometry>,
    msg: PoseStamped,
    secs: i32,
    nsecs: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    // Create shared timestamp
    let timestamp = r2r::builtin_interfaces::msg::Time {
        sec: secs,
        nanosec: nsecs,
    };
    
    // Create and publish TF message
    let tf_transform = TransformStamped {
        header: r2r::std_msgs::msg::Header {
            stamp: timestamp.clone(),
            frame_id: ODOM_FRAME.to_string(),
        },
        child_frame_id: BASE_LINK_FRAME.to_string(),
        transform: r2r::geometry_msgs::msg::Transform {
            translation: r2r::geometry_msgs::msg::Vector3 {
                x: msg.pose.position.x,
                y: msg.pose.position.y,
                z: msg.pose.position.z,
            },
            rotation: msg.pose.orientation.clone(),
        },
    };
    
    let pose = msg.pose;
    
    let tf_msg = TFMessage {
        transforms: vec![tf_transform],
    };
    
    tf_puber.publish(&tf_msg)
        .map_err(|e| format!("Failed to publish TF: {}", e))?;
    
    // Create and publish odometry message
    let odom_msg = Odometry {
        header: r2r::std_msgs::msg::Header {
            stamp: timestamp,
            frame_id: ODOM_FRAME.to_string(),
        },
        child_frame_id: BASE_LINK_FRAME.to_string(),
        pose: r2r::geometry_msgs::msg::PoseWithCovariance {
            pose,
            covariance: vec![0.0; 36], // Default covariance
        },
        twist: r2r::geometry_msgs::msg::TwistWithCovariance::default(), // Zero velocity
    };
    
    odom_puber.publish(&odom_msg)
        .map_err(|e| format!("Failed to publish odometry: {}", e).into())
}

/// Initialize and run the go2 base node
fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Setup Ctrl+C handler
    ctrlc::set_handler(move || {
        if RUNNING.load(Ordering::SeqCst) {
            println!("\nReceived Ctrl+C, shutting down...");
            RUNNING.store(false, Ordering::SeqCst);
        }
    })?;

    // Initialize ROS2 context and node
    let ctx = r2r::Context::create()?;
    let mut node = r2r::Node::create(ctx, NODE_NAME, "")?;
    
    // Create publishers
    let req_puber = node.create_publisher::<Request>(SPORT_REQUEST_TOPIC, QosProfile::default())?;
    let odom_puber = node.create_publisher::<Odometry>(ODOM_TOPIC, QosProfile::default())?;
    let tf_puber = node.create_publisher::<TFMessage>(TF_TOPIC, QosProfile::default())?;
    
    // Create subscribers
    let cmd_vel_subscriber = node.subscribe::<Twist>(CMD_VEL_TOPIC, QosProfile::default())?;
    let pose_sub = node.subscribe::<PoseStamped>(POSE_TOPIC, QosProfile::default())?;
    
    // Print startup information
    println!("🤖 {} node started successfully", NODE_NAME);
    println!("📡 Subscribing to: {} and {}", CMD_VEL_TOPIC, POSE_TOPIC);
    println!("📤 Publishing to: {}, {} and {}", SPORT_REQUEST_TOPIC, ODOM_TOPIC, TF_TOPIC);
    println!("🔄 Ready to process messages...");
    
    // Setup async task execution
    let mut pool = LocalPool::new();
    let spawner = pool.spawner();
    
    // Clone publishers for async tasks
    let odom_puber_clone = odom_puber.clone();
    let tf_puber_clone = tf_puber.clone();
    
    // Spawn cmd_vel handling task
    let req_puber_clone = req_puber.clone();
    spawner.spawn_local(async move {
        let mut sub = cmd_vel_subscriber;
        while let Some(msg) = sub.next().await {
            if !RUNNING.load(Ordering::SeqCst) {
                break;
            }
            handle_cmd_vel(&msg, &req_puber_clone);
        }
    })?;
    
    // Spawn pose handling task
    spawner.spawn_local(async move {
        pose_sub.for_each(|msg| {
            if !RUNNING.load(Ordering::SeqCst) {
                return futures::future::ready(());
            }
            if let Err(e) = handle_pose(msg, &odom_puber_clone, &tf_puber_clone) {
                eprintln!("❌ Error handling pose message: {}", e);
            }
            futures::future::ready(())
        }).await;
    }).map_err(|e| format!("Failed to spawn pose handling task: {}", e))?;
    
    // Main event loop
    const SPIN_DURATION_MS: u64 = 10;
    let spin_duration = std::time::Duration::from_millis(SPIN_DURATION_MS);
    
    println!("🚀 Starting main event loop (spin interval: {}ms)", SPIN_DURATION_MS);
    
    while RUNNING.load(Ordering::SeqCst) {
        node.spin_once(spin_duration);
        pool.run_until_stalled();
        std::thread::sleep(std::time::Duration::from_millis(10));
    }
    
    // Send stop command before exit
    send_stop_command(&req_puber);
    println!("Node stopped.");
    
    Ok(())
}